from langchain_openai import ChatOpenAI
import os 
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import HumanMessage,AIMessage,SystemMessage
import json
import re
from typing import List
llm = ChatOpenAI(
  model = 'deepseek-chat',
  base_url='https://api.deepseek.com/v1',
  api_key=os.environ['DEEPSEEK_API_KEY']
)
prompt = ChatPromptTemplate.from_messages(
  [('system','Answer the user query. Wrap the output in `json`'),
('human','{input}'),]
)

def extract_json(AI_message: AIMessage) -> List:
  """
  从AIMessage中提取json字符串
  """
  try:
    content = AI_message.content
    # re = r'\`\`\`json(.*?)\`\`\`'
    pattern = r"\`\`\`json(.*?)\`\`\`"
    matches = re.findall(pattern,content,re.DOTALL)
    return [json.loads(match) for match in matches]
  except Exception as e:
    raise ValueError(f'json提取失败:{content} {e}')
chain = prompt|llm|extract_json
print(chain.invoke({'input':'我叫姚明，今年18岁，来自武汉，手机号是1231412，邮箱地址是1231412@qq.com'}))
